import numpy as np
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure
from matplotlib.patches import Rectangle
from PyQt5.QtWidgets import QApplication, QMainWindow
import sys


class SelectBox:
    def __init__(self, ax, data_points):
        """
        初始化选择框工具。
        """
        self.ax = ax
        self.data_points = data_points  # 储存原始数据点
        self.highlighted_points = None  # 存储高亮的数据点
        self.rect = None  # 用于存储矩形选择框的对象
        self.press = None  # 记录鼠标按下时的位置

    def highlight_selected_points(self, x0, y0, x1, y1):
        """根据选择的区域高亮点"""
        if x0 > x1:
            temp = x0
            x0 = x1
            x1 = temp
        if y0 > y1:
            temp = y0
            y0 = y1
            y1 = temp
        in_rect = (self.data_points[:, 0] >= x0) & (self.data_points[:, 0] <= x1) & \
                (self.data_points[:, 1] >= y0) & (self.data_points[:, 1] <= y1)
        self.highlighted_points = self.data_points[in_rect]
        self.update_plot()

    def update_plot(self):
        """更新图表以高亮显示点"""
        self.ax.clear()  # 先清除原有图形
        self.ax.plot(self.data_points[:, 0], self.data_points[:, 1], 'o', color='blue')  # 绘制所有点
        if self.highlighted_points is not None:
            self.ax.plot(self.highlighted_points[:, 0], self.highlighted_points[:, 1], 'o', color='red',
                         markersize=10)  # 高亮点
        if self.rect is not None:
            self.rect.set_visible(True)  # 保持矩形可见
        self.ax.figure.canvas.draw_idle()

    def on_press(self, event):
        """
        鼠标按下事件处理器：记录按下位置。
        """
        if event.button != 1:  # 确保是左键按下
            return
        # 这行代码检查是否是鼠标左键被按下。如果不是（即用户点击的是右键或其他键），则直接返回，不做任何处理。这是因为这个方法主要关注对左键操作的响应。
        if event.inaxes != self.ax:
            return
        # 这行代码检查鼠标按下事件是否发生在当前关心的坐标轴 (self.ax) 上。如果不在，也直接返回。这样可以确保只处理与当前图形相关的交互事件。
        self.press = (event.xdata, event.ydata)  # 记录按下时的坐标

    def on_release(self, event):
        """
        鼠标释放事件处理器：完成选择并显示选择区域。
        """
        if event.button != 1 or self.press is None:  # 确保是左键且有按下记录
            return
        if event.inaxes != self.ax:
            return
        x0, y0 = self.press
        x1, y1 = event.xdata, event.ydata
        # 计算并打印选择的区域
        print(f"Selected area: ({x0:.2f}, {y0:.2f}) to ({x1:.2f}, {y1:.2f})")
        # 清理并重绘图形，移除之前的矩形准备下次选择
        if self.rect is not None:
            self.rect.remove()
            self.rect = None
        self.ax.figure.canvas.draw_idle()
        if self.press is not None:
            x0, y0 = self.press
            x1, y1 = event.xdata, event.ydata
            self.highlight_selected_points(x0, y0, x1, y1)

    def on_motion(self, event):
        """
        鼠标移动事件处理器：实时更新选择框。
        """
        if self.press is None or event.button != 1:  # 确保左键被按下的情况下处理移动
            return
        if event.inaxes != self.ax:
            return

        x0, y0 = self.press
        x1, y1 = event.xdata, event.ydata

        # 更新矩形的坐标保证是矩形
        x0, x1 = sorted([x0, x1])
        y0, y1 = sorted([y0, y1])
        # 更新或创建矩形框
        if self.rect is not None:
            self.rect.set_xy((x0, y0))
            self.rect.set_width(x1 - x0)
            self.rect.set_height(y1 - y0)
        else:
            self.rect = Rectangle((x0, y0), x1 - x0, y1 - y0, edgecolor='r', facecolor='none', linewidth=1)
            self.ax.add_patch(self.rect)

        self.ax.figure.canvas.draw_idle()


class DualAxisWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.initUI()

    def initUI(self):
        """
        初始化用户界面。
        """
        self.fig = Figure()
        self.ax = self.fig.add_subplot(111)
        np.random.seed(0)  # 为了可重复性设置随机种子
        self.data_points = np.random.rand(10, 2)  # 假设有两个维度的数据点
        self.select_box = SelectBox(self.ax, self.data_points)
        self.canvas = FigureCanvas(self.fig)
        self.setCentralWidget(self.canvas)
        self.canvas.mpl_connect('button_press_event', self.select_box.on_press)
        # 这行代码将鼠标按钮按下事件与SelectBox类中的on_press方法关联起来。
        # 也就是说，当用户在图形界面上按下鼠标按钮时，on_press方法会被触发执行。
        # 在这个方法中，通常会记录鼠标按下时的位置，作为后续操作（如选择区域）的起始坐标。
        self.canvas.mpl_connect('button_release_event', self.select_box.on_release)
        self.canvas.mpl_connect('motion_notify_event', self.select_box.on_motion)
        self.select_box.update_plot()  # 初始绘制图表


if __name__ == '__main__':
    app = QApplication([])
    mainWin = DualAxisWindow()
    mainWin.show()
    sys.exit(app.exec_())